import numpy as np
import heapq

# 1. 生成随机向量数据
def generate_vectors(num_vectors=20, dimensions=2):
    """ 
    生成向量数据
    :param num_vectors: 向量数量
    :param dimensions: 向量维度
    :return: 随机生成的向量数据库
    """ 
    np.random.seed(42)
    return np.random.rand(num_vectors, dimensions)

data_vectors = generate_vectors(num_vectors=20, dimensions=2)

# 2. 构建分层索引
class HNSWLayer:
    """ 
    表示 HNSW 简单层图结构
    """ 
    def __init__(self, layer_id, vectors):
        self.layer_id = layer_id
        self.vectors = vectors
        self.graph = {i: [] for i in range(len(vectors))}
    
    def add_edge(self, node_a, node_b):
        self.graph[node_a].append(node_b)
        self.graph[node_b].append(node_a)

def construct_hnsw(vectors, max_layers=3, neighbors_per_node=3):
    """ 
    构造 HNSW 分层索引
    :param vectors: 数据向量
    :param max_layers: 最大层数
    :param neighbors_per_node: 每个节点的近邻数
    :return: 分层索引
    """ 
    layers = []
    current_vectors = vectors
    for layer_id in range(max_layers):
        layer = HNSWLayer(layer_id, current_vectors)
        num_vectors = len(current_vectors)
        
        for i in range(num_vectors):
            distances = []
            for j in range(num_vectors):
                if i != j:
                    dist = np.linalg.norm(
                        current_vectors[i] - current_vectors[j])
                    distances.append((dist, j))
            
            # 选出最近的 neighbors_per_node 个邻居
            nearest_neighbors = heapq.nsmallest(neighbors_per_node, distances)
            for _, neighbor_idx in nearest_neighbors:
                layer.add_edge(i, neighbor_idx)
        
        layers.append(layer)
        # 上一层的点减少为当前点的一半
        current_vectors = current_vectors[::2]
    return layers

hnsw_layers = construct_hnsw(data_vectors, max_layers=3, neighbors_per_node=3)

# 3. 分层搜索算法
def hierarchical_search(query_vector, hnsw_layers, start_layer=2):
    """ 
    在分层索引中执行分层搜索
    :param query_vector: 查询向量
    :param hnsw_layers: 分层索引
    :param start_layer: 开始的层数
    :return: 底层找到的最近邻索引和距离
    """ 
    current_node = 0  # 假设从第 0 节点开始
    for layer in range(start_layer, -1, -1):
        graph = hnsw_layers[layer].graph
        vectors = hnsw_layers[layer].vectors
        min_distance = np.linalg.norm(query_vector - vectors[current_node])
        nearest_node = current_node
        
        while True:
            updated = False
            for neighbor in graph[current_node]:
                distance = np.linalg.norm(query_vector - vectors[neighbor])
                if distance < min_distance:
                    min_distance = distance
                    nearest_node = neighbor
                    updated = True
            
            if not updated:
                break
            current_node = nearest_node
    
    return current_node, min_distance

# 4. 测试分层搜索
query_vector = np.array([0.5, 0.5])  # 查询向量
nearest_node, distance = hierarchical_search(
    query_vector, hnsw_layers, start_layer=2)

# 5. 输出结果
print("向量数据:")
print(data_vectors)
print("\nHNSW 分层索引结构:")
for layer in hnsw_layers:
    print(f"层 {layer.layer_id}: {layer.graph}")

print("\n查询向量:", query_vector)
print(f"最近邻节点索引: {nearest_node}, 距离: {distance:.6f}")